# main.py
# import os
# import argparse
# from dotenv import load_dotenv
#
# # rag、utils、documents包都与main.py在同一目录，可以直接导入
# from rag.rag_system import RAGSystem
# from utils.file_utils import print_system_info
#
# def setup_environment():
#     """设置环境"""
#     load_dotenv()
#     if not os.getenv('DEEPSEEK_API_KEY'):
#         print("❌ 未找到DEEPSEEK_API_KEY环境变量")
#         print("请在.env文件中设置: DEEPSEEK_API_KEY=你的API密钥")
#         return False
#     return True
#
# def build_knowledge_base():
#     """构建知识库"""
#     print("🔨 构建知识库...")
#     rag_system = RAGSystem()
#
#     # 【关键修正】更新文档路径，因为它现在在 documents 包里
#     doc_path = "documents/sample_document.txt"
#
#     try:
#         documents = rag_system.load_documents(doc_path)
#         chunks = rag_system.split_documents(documents)
#         vector_store_path = "vector_store/ai_knowledge_base"
#         rag_system.create_vector_store(chunks, vector_store_path)
#         rag_system.create_qa_chain()
#         print("✅ 知识库构建完成！")
#         return rag_system
#     except Exception as e:
#         print(f"❌ 构建知识库失败: {e}")
#         return None
#
# def load_existing_knowledge_base():
#     """加载已有知识库"""
#     print("📂 加载已有知识库...")
#     vector_store_path = "vector_store/ai_knowledge_base"
#     if not os.path.exists(vector_store_path):
#         print(f"❌ 向量存储不存在: {vector_store_path}")
#         return None
#     try:
#         rag_system = RAGSystem()
#         rag_system.load_vector_store(vector_store_path)
#         rag_system.create_qa_chain()
#         print("✅ 知识库加载完成！")
#         return rag_system
#     except Exception as e:
#         print(f"❌ 加载知识库失败: {e}")
#         return None
#
# # ... (interactive_qa, test_qa_system, main 函数保持不变) ...
# def interactive_qa(rag_system):
#     """交互式问答"""
#     print("\n" + "=" * 60)
#     print("🤖 单文档问答系统")
#     print("=" * 60)
#     print("输入 '退出' 或 'quit' 结束对话")
#     print("输入 '搜索 <关键词>' 搜索相关文档片段")
#     print("输入 '重建' 重新构建知识库")
#     print("-" * 60)
#
#     while True:
#         try:
#             user_input = input("\n💬 你的问题: ").strip()
#
#             if not user_input:
#                 continue
#
#             if user_input.lower() in ['退出', 'quit', 'exit']:
#                 print("👋 再见！")
#                 break
#
#             elif user_input.lower() == '重建':
#                 print("🔄 重新构建知识库...")
#                 rag_system = build_knowledge_base()
#                 if rag_system:
#                     continue
#                 else:
#                     break
#
#             elif user_input.startswith('搜索 '):
#                 query = user_input[3:].strip()
#                 if query:
#                     print(f"🔍 搜索: {query}")
#                     results = rag_system.search_similar_documents(query, k=2)
#
#                     if results:
#                         for result in results:
#                             print(f"\n📄 相关文档 {result['rank']}:")
#                             print(f"内容: {result['content'][:200]}...")
#                             if result['metadata']:
#                                 print(f"元数据: {result['metadata']}")
#                     else:
#                         print("❌ 未找到相关文档")
#                 continue
#
#             # 普通问答
#             print("\n🤔 思考中...", end="", flush=True)
#             answer = rag_system.ask_question(user_input)
#             print(f"\r✅ 答案: {answer}")
#
#         except KeyboardInterrupt:
#             print("\n\n🛑 程序被中断")
#             break
#         except Exception as e:
#             print(f"\n❌ 错误: {e}")
#
# def test_qa_system(rag_system):
#     """测试问答系统"""
#     print("\n🧪 系统测试...")
#
#     test_questions = [
#         "人工智能是什么时候提出的？",
#         "什么是深度学习？",
#         "列举几个AI寒冬的原因",
#         "GPT-3是什么时候发布的？",
#         "人工智能的未来发展趋势有哪些？"
#     ]
#
#     for i, question in enumerate(test_questions, 1):
#         print(f"\n{i}. 问题: {question}")
#         answer = rag_system.ask_question(question)
#         print(f"   答案: {answer}")
#
#         # 显示相关文档
#         results = rag_system.search_similar_documents(question, k=1)
#         if results:
#             print(f"   来源: {results[0]['content'][:100]}...")
#
# def main():
#     """主函数"""
#     print_system_info()
#     if not setup_environment():
#         return
#
#     parser = argparse.ArgumentParser(description="单文档问答系统")
#     parser.add_argument('--build', action='store_true', help='重新构建知识库')
#     parser.add_argument('--test', action='store_true', help='运行测试')
#
#     args = parser.parse_args()
#
#     if args.build or not os.path.exists("vector_store/ai_knowledge_base"):
#         rag_system = build_knowledge_base()
#     else:
#         rag_system = load_existing_knowledge_base()
#
#     if not rag_system:
#         print("❌ 无法初始化RAG系统")
#         return
#
#     if args.test:
#         test_qa_system(rag_system)
#     else:
#         interactive_qa(rag_system)
#
# if __name__ == "__main__":
#     main()

# 以上是出来txt的代码

# main.py
import os
import sys
from dotenv import load_dotenv
# 导入我们的核心模块
from rag.rag_system import RAGSystem
from utils.file_utils import print_system_info


def get_vector_store_path(file_path: str) -> str:
    """
    根据文档文件路径，生成唯一的向量存储路径。
    例如: documents/report.pdf -> vector_store/report_faiss
    """
    # 获取不含扩展名的文件名
    base_name = os.path.splitext(os.path.basename(file_path))[0]
    # 存储在vector_store目录下，文件名后加_faiss以区分
    vector_store_path = os.path.join("vector_store", f"{base_name}_faiss")
    return vector_store_path


def handle_load_command(file_path: str, current_rag_system) -> tuple[RAGSystem, str]:
    """
    处理‘加载’命令：如果知识库存在则加载，不存在则构建。
    返回加载好的RAGSystem实例和向量库路径。
    """
    print(f"🔍 正在处理文件: {file_path}")
    # 1. 检查文件是否存在和是否是支持的类型
    from utils.file_utils import validate_document
    from utils.config import SUPPORTED_FORMATS
    file_extension = os.path.splitext(file_path)[1].lower()
    if not os.path.exists(file_path):
        print(f"❌ 错误：文件 '{file_path}' 不存在。请检查路径。")
        return current_rag_system, None
    if file_extension not in SUPPORTED_FORMATS:
        print(f"❌ 错误：不支持的文件格式 '{file_extension}'。支持的格式: {SUPPORTED_FORMATS}")
        return current_rag_system, None
    # 2. 确定向量存储路径
    vector_store_path = get_vector_store_path(file_path)
    # 3. 判断是加载还是构建
    if os.path.exists(vector_store_path):
        print(f"📂 找到已存在的知识库，正在加载...")
        try:
            rag_system = RAGSystem()
            rag_system.load_vector_store(vector_store_path)
            rag_system.create_qa_chain()
            print(f"✅ 知识库加载成功！来源: {file_path}")
            return rag_system, vector_store_path
        except Exception as e:
            print(f"❌ 加载知识库失败: {e}")
            return current_rag_system, None
    else:
        print(f"🔨 未找到知识库，正在为你构建...")
        try:
            rag_system = RAGSystem()
            documents = rag_system.load_documents(file_path)
            chunks = rag_system.split_documents(documents)
            rag_system.create_vector_store(chunks, vector_store_path)
            rag_system.create_qa_chain()
            print(f"✅ 知识库构建并加载成功！来源: {file_path}")
            return rag_system, vector_store_path
        except Exception as e:
            print(f"❌ 构建知识库失败: {e}")
            return current_rag_system, None


def interactive_mode():
    """主交互循环"""
    current_rag_system = None
    current_vector_store_path = None
    print_system_info()
    if not load_dotenv() or not os.getenv('DEEPSEEK_API_KEY'):
        print("❌ 未找到DEEPSEEK_API_KEY，程序无法运行。请检查 .env 文件。")
        return
    print("\n" + "=" * 60)
    print("🤖 欢迎使用外挂知识库问答系统！")
    print("=" * 60)
    print("可用命令:")
    print("  加载 <文件路径>      - 指定一个.txt或.pdf文件来构建或加载知识库")
    print("  比如：加载 documents/sample_document.txt 或 加载 documents/my_report.pdf")
    print("  退出 或 quit         - 结束程序")
    print("-" * 60)
    while True:
        try:
            # 提示用户输入
            prompt = "💬 你的命令或问题: "
            if not current_rag_system:
                prompt = "💬 请先使用 '加载 <文件路径>' 来指定知识库: "
            user_input = input(prompt).strip()
            if not user_input:
                continue
            # 处理退出命令
            if user_input.lower() in ['退出', 'quit', 'exit']:
                print("👋 再见！")
                break
            # 处理加载命令
            if user_input.startswith('加载 '):
                file_path = user_input[3:].strip().strip('"\'')  # 去除多余的引号
                current_rag_system, current_vector_store_path = handle_load_command(file_path, current_rag_system)
                continue
            # 如果没有加载知识库，则无法回答问题
            if not current_rag_system:
                print("💡 提示：请先使用 '加载 <文件路径>' 来指定一个知识库。")
                continue
            # 普通问答
            print("\n🤔 思考中...", end="", flush=True)
            answer = current_rag_system.ask_question(user_input)
            print(f"\r✅ 答案: {answer}")
        except KeyboardInterrupt:
            print("\n\n🛑 程序被中断")
            break
        except Exception as e:
            print(f"\n❌ 发生未知错误: {e}")


if __name__ == "__main__":
    interactive_mode()